from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD

MEAN_ZERO = (0.0, 0.0, 0.0)
STA_ONE = (1.0, 1.0, 1.0)


def get_url(ckpt_name):
    url = "https://ml.cs.tsinghua.edu.cn/~xiaoyang/aresbench/ckpt-imagenet/" + ckpt_name

    return url


imagenet_model_zoo = {
    "resnet50_normal": {
        "model": "tv_resnet50",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "resnet50_moco": {
        "model": "tv_resnet50",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ResNet50_MOCO.pth"),
        "pt": "ARES_ResNet50_MOCO.pth"
    },
    "resnet50_at": {
        "model": "resnet50",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ResNet50_AT.pth"),
        "pt": "ARES_ResNet50_AT.pth"
    },
    "resnet50_sota1": {
        "model": "resnet50_rl",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ResNet50_RB1.pth"),
        "pt": "ARES_ResNet50_RB1.pth"
    },
    "resnet50_sota2": {
        "model": "resnet50_rl",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ResNet50_RB2.pth"),
        "pt": "ARES_ResNet50_RB2.pth"
    },
    "resnet101_normal": {
        "model": "tv_resnet101",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "resnet101_at": {
        "model": "resnet101",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ResNet101_AT.pth"),
        "pt": "ARES_ResNet101_AT.pth"
    },
    "resnet152_normal": {
        "model": "tv_resnet152",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 5,
        "url": "",
    },
    "resnet152_at": {
        "model": "resnet152",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": True,
        "batch_size": 5,
        "url": get_url("ARES_ResNet152_AT.pth"),
        "pt": "ARES_ResNet152_AT.pth"
    },
    "resnet152_fd": {
        "model": "resnet152_fd",
        "mean": MEAN_ZERO,
        "std": STA_ONE,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 5,
        "url": get_url("ARES_ResNet152_FD.pth"),
        "pt": "ARES_ResNet152_FD.pth"
    },
    "wresnet50_normal": {
        "model": "wide_resnet50_2",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "wresnet50_at": {
        "model": "wide_resnet50_2",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": True,
        "batch_size": 25,
        "url": get_url("ARES_Wide_ResNet50_2_AT.pth"),
        "pt": "ARES_Wide_ResNet50_2_AT.pth"
    },
    "wresnet50_sota": {
        "model": "wide_resnet50_2_rl",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Wide_ResNet50_2_RB.pth"),
        "pt": "ARES_Wide_ResNet50_2_RB.pth"
    },
    "vgg13_normal": {
        "model": "vgg13",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "vgg16_normal": {
        "model": "vgg16",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "vgg19_normal": {
        "model": "vgg19",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "densenet121_normal": {
        "model": "tv_densenet121",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "densenet161_normal": {
        "model": "densenet161",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "densenet201_normal": {
        "model": "densenet201",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "convs_normal": {
        "model": "convnext_small",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "convs_21k": {
        "model": "convnext_small_in22ft1k",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "convb_normal": {
        "model": "convnext_base",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "convb_21k": {
        "model": "convnext_base_in22ft1k",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "convl_normal": {
        "model": "convnext_large",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 5,
        "url": "",
    },
    "convl_21k": {
        "model": "convnext_large_in22ft1k",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 5,
        "url": "",
    },
    "swins_normal": {
        "model": "swin_small_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "swins_21k": {
        "model": "swin_small_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Swin_small_patch4_window7_224_21K.pth"),
        "pt": "ARES_Swin_small_patch4_window7_224_21K.pth"
    },
    "swins_at": {
        "model": "swin_small_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Swin_small_patch4_window7_224_AT.pth"),
        "pt": "ARES_Swin_small_patch4_window7_224_AT.pth"
    },
    "swinb_normal": {
        "model": "swin_base_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Swin_base_patch4_window7_224_Normal.pth"),
        "pt": "ARES_Swin_base_patch4_window7_224_Normal.pth"
    },
    "swinb_21k": {
        "model": "swin_base_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "swinb_at": {
        "model": "swin_base_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Swin_base_patch4_window7_224_AT.pth"),
        "pt": "ARES_Swin_base_patch4_window7_224_AT.pth"
    },
    "swinl_21k": {
        "model": "swin_large_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 5,
        "url": "",
    },
    "swinl_at": {
        "model": "swin_large_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 5,
        "url": get_url("ARES_Swin_large_patch4_window7_224_AT.pth"),
        "pt": "ARES_Swin_large_patch4_window7_224_AT.pth"
    },
    "vits_normal": {
        "model": "vit_small_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ViT_small_patch16_224_Normal.pth"),
        "pt": "ARES_ViT_small_patch16_224_Normal.pth"
    },
    "vits_21k": {
        "model": "vit_small_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "vits_at": {
        "model": "vit_small_patch16_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ViT_small_patch16_224_AT.pth"),
        "pt": "ARES_ViT_small_patch16_224_AT.pth"
    },
    "vitb_normal": {
        "model": "vit_base_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ViT_base_patch16_224_Normal.pth"),
        "pt": "ARES_ViT_base_patch16_224_Normal.pth"
    },
    "vitb_21k": {
        "model": "vit_base_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "vitb_at": {
        "model": "vit_base_patch16_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ViT_base_patch16_224_AT.pth"),
        "pt": "ARES_ViT_base_patch16_224_AT.pth"
    },
    "vitb_mae": {
        "model": "vit_base_patch16",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ViT_base_patch16_224_MAE.pth"),
        "pt": "ARES_ViT_base_patch16_224_MAE.pth"
    },
    "vitl_normal": {
        "model": "vit_large_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 5,
        "url": get_url("ARES_ViT_large_patch16_224_Normal.pth"),
        "pt": "ARES_ViT_large_patch16_224_Normal.pth"
    },
    "vitl_21k": {
        "model": "vit_large_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 5,
        "url": "",
    },
    "vitl_mae": {
        "model": "vit_large_patch16",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 5,
        "url": get_url("ARES_ViT_large_patch16_224_MAE.pth"),
        "pt": "ARES_ViT_large_patch16_224_MAE.pth"
    },
    "t2t14_normal": {
        "model": "t2t_vit_t_14",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_T2T_ViT_14_Normal.pth"),
        "pt": "ARES_T2T_ViT_14_Normal.pth"
    },
    "t2t19_normal": {
        "model": "t2t_vit_t_19",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_T2T_ViT_19_Normal.pth"),
        "pt": "ARES_T2T_ViT_19_Normal.pth"
    },
    "t2t24_normal": {
        "model": "t2t_vit_t_24",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_T2T_ViT_24_Normal.pth"),
        "pt": "ARES_T2T_ViT_24_Normal.pth"
    },
    "xcits_normal": {
        "model": "xcit_small_24_p16_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "xcits_sota": {
        "model": "xcit_small_12_p16_224",
        "mean": MEAN_ZERO,
        "std": STA_ONE,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_XciT_small_12_p16_224_RB.pth"),
        "pt": "ARES_XciT_small_12_p16_224_RB.pth"
    },
    "xcitm_normal": {
        "model": "xcit_medium_24_p16_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "xcitm_sota": {
        "model": "xcit_medium_12_p16_224",
        "mean": MEAN_ZERO,
        "std": STA_ONE,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_XciT_medium_12_p16_224_RB.pth"),
        "pt": "ARES_XciT_medium_12_p16_224_RB.pth"
    },
    "xcitl_normal": {
        "model": "xcit_large_24_p16_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "xcitl_sota": {
        "model": "xcit_large_12_p16_224",
        "mean": MEAN_ZERO,
        "std": STA_ONE,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_XciT_large_12_p16_224_RB.pth"),
        "pt": "ARES_XciT_large_12_p16_224_RB.pth"
    },
    "convnexts_at": {
        "model": "convnext_small",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ConvNext_Small_AT.pth"),
        "pt": "ARES_ConvNext_Small_AT.pth"
    },
    "convnextb_at": {
        "model": "convnext_base",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ConvNext_Base_AT.pth"),
        "pt": "ARES_ConvNext_Base_AT.pth"
    },
    "convnextl_at": {
        "model": "convnext_large",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ConvNext_Large_AT.pth"),
        "pt": "ARES_ConvNext_Large_AT.pth"
    },
}
